#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use 
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

import torch
import math
from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer
from scene.gaussian_model import GaussianModel
from utils.sh_utils import eval_sh

def render(viewpoint_camera, 
           pc : GaussianModel, 
           pipe, 
           bg_color : torch.Tensor, 
           scaling_modifier = 1.0, 
           override_color = None):
    """
    Render the scene. 
    
    Background tensor (bg_color) must be on GPU!
    """
 
    # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
    screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0
    try:
        screenspace_points.retain_grad()
    except:
        pass

    # Set up rasterization configuration
    # tanfovx 是水平视角一半的正切值，它定义了视锥（view frustum）在水平方向的范围
    #                           可以用来确定屏幕上每个像素在水平位置上世界空间单位
    # tanfovy 是垂直视角一半的正切值，它定义了视锥（view frustum）在垂直方向的范围
    #                           可以用来确定屏幕上每个像素在垂直位置上世界空间单位
    tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)
    tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)

    raster_settings = GaussianRasterizationSettings(
        image_height=int(viewpoint_camera.image_height),
        image_width=int(viewpoint_camera.image_width),
        tanfovx=tanfovx,
        tanfovy=tanfovy,
        bg=bg_color,
        scale_modifier=scaling_modifier,
        viewmatrix=viewpoint_camera.world_view_transform,
        projmatrix=viewpoint_camera.full_proj_transform,
        sh_degree=pc.active_sh_degree,
        campos=viewpoint_camera.camera_center,
        prefiltered=False,
        debug=pipe.debug
    )

    rasterizer = GaussianRasterizer(raster_settings=raster_settings)

    means3D = pc.get_xyz
    means2D = screenspace_points
    opacity = pc.get_opacity
    scales = pc.get_scaling
    rotations = pc.get_rotation
        
    visible_mask = rasterizer.markVisible(positions=means3D)
    colors = pc.feature_encoder(viewpoint_camera, visible_mask)

    # Rasterize visible Gaussians to image, obtain their radii (on screen). 
    rendered_image, radii = rasterizer(means3D=means3D[visible_mask],
                                       means2D=means2D[visible_mask],
                                       shs=None,
                                       colors_precomp=colors,
                                       opacities=opacity[visible_mask],
                                       scales=scales[visible_mask],
                                       rotations=rotations[visible_mask],
                                       cov3D_precomp=None)

    tem_radii = torch.zeros_like(visible_mask, dtype=radii.dtype, device=visible_mask.device)
    tem_radii[visible_mask] = radii
    visibility_filter = torch.logical_and(visible_mask, tem_radii>0)
    
    return {"render": rendered_image,
            "viewspace_points": screenspace_points,
            "visibility_filter" : visibility_filter,
            "radii": tem_radii}
